package middlewares

import (
	"github.com/didip/tollbooth"
	"github.com/didip/tollbooth/limiter"
	"github.com/gin-gonic/gin"
)

func LimiterMiddleware(limit *limiter.Limiter) gin.HandlerFunc {
	return func(context *gin.Context) {
		limit.SetMessage("服务繁忙，请稍后再试...")
		httpError := tollbooth.LimitByRequest(limit, context.Writer, context.Request)
		if httpError != nil {
			context.Data(httpError.StatusCode, limit.GetMessageContentType(), []byte(httpError.Message))
			context.Abort()
		} else {
			context.Next()
		}
	}
}
